import pandas as pd
import numpy as np
import networkx as nx
import gymnasium
from gymnasium import spaces

route_position_file = './Data/data_50.csv'
position_distance_file = './Data/distance_50.csv'
charging_info_file = './Data/roads_50.csv'
speed_info_file = './Data/speed_50.csv'
car_info_file = './Data/EVs_50.csv'


class RouteInfo():
    def __init__(self) -> None:
        # 1. 读取巴士站点地理位置数据
        self.stations = pd.read_csv(route_position_file, header=None, names=['Longitude', 'Latitude'])
        self.stations['StationID'] = self.stations.index
        # 2. 读取站点间距离数据
        self.distance_matrix = pd.read_csv(position_distance_file, header=None)
        # 3. 读取路段充电信息
        self.charging_matrix = pd.read_csv(charging_info_file, header=None)
        # 4. 读取路段速度信息
        self.speed_matrix = pd.read_csv(speed_info_file, header=None)
        # 5. 读取电动汽车信息
        column_names = ['StartLongitude', 'StartLatitude', 'EndLongitude', 'EndLatitude',
                    'InitialEnergy', 'BatteryCapacity', 'Deadline', 'ConsumptionRate']
        self.car_matrix = pd.read_csv(car_info_file, header=None, names=column_names)
        # 匹配电动汽车的起点和终点站点ID
        self.car_matrix['StartStation'] = self.car_matrix.apply(lambda row: self.__find_nearest_station(row['StartLongitude'], row['StartLatitude'], self.stations), axis=1)
        self.car_matrix['EndStation'] = self.car_matrix.apply(lambda row: self.__find_nearest_station(row['EndLongitude'], row['EndLatitude'], self.stations), axis=1)

        # 道路无向图
        self.__build_graph()

    def __find_nearest_station(self, longitude, latitude, stations):
        stations['Distance'] = np.sqrt((stations['Longitude'] - longitude)**2 + (stations['Latitude'] - latitude)**2)
        nearest_station = stations.loc[stations['Distance'].idxmin()]
        return int(nearest_station['StationID'])

    def __build_graph(self):
        self.graph = nx.Graph()
        for index, row in self.stations.iterrows():
            self.graph.add_node(int(row['StationID']), pos=(row['Longitude'], row['Latitude']))
        num_stations = len(self.stations)
        edges = []
        for start in range(num_stations):
            for end in range(start + 1, num_stations):
                distance = self.distance_matrix.iloc[start, end]
                is_charging = self.charging_matrix.iloc[start, end]
                speed = self.speed_matrix.iloc[start, end]
                time = distance / speed  # 时间 = 距离 / 速度

                edges.append({
                    'StartStation': start,
                    'EndStation': end,
                    'Distance': distance,
                    'IsCharging': is_charging,
                    'Speed': speed,
                    'Time': time
                })

        for edge in edges:
            self.graph.add_edge(edge['StartStation'], edge['EndStation'],
                                Distance = edge['Distance'], IsCharging = edge['IsCharging'],
                                Speed = edge['Speed'], Time = edge['Time'])


class RouteEnv(gymnasium.Env):
    def __init__(self, car_id) -> None:
        super(RouteEnv, self).__init__()
        
        self.route_info = RouteInfo()
        self.grid_size = 50
        self.car_num = 50
        self.car_index = car_id
        self.charging_power = 100    # 充电速率 KWh
        self.dead_line_time = 0
        self.path = []
        
        self.action_low = 0
        self.action_high = 49
        self.action_space = spaces.Discrete(50)
        self.observation_space = spaces.Box(low=np.array([0, 0, 0]),                     # [current_position.index, time, energy]
                                             high=np.array([self.grid_size-1, np.inf, 100]), 
                                             dtype=np.float32)
        self.reset(car_id)

    def reset(self, car_index=None, seed=None):
        if not car_index == None:
            car_info = self.route_info.car_matrix.loc[car_index, :]
        else:
            car_info = self.route_info.car_matrix.loc[self.car_index, :]
        self.current_position = car_info['StartStation']
        self.end_position = car_info['EndStation']
        self.energy = car_info['InitialEnergy']
        self.energy_history = self.energy
        self.battery_capacity = car_info['BatteryCapacity']
        self.time = car_info['Deadline']
        self.dead_line_time = car_info['Deadline']
        self.consumption_rate = car_info['ConsumptionRate']
        self.path.clear()
        self.path.append(self.current_position)

        return self.get_observation(), {}

    def get_reward(self):
        done = False
        resolved = False
        reward = 0

        if self.energy < 0\
                or self.energy > self.battery_capacity\
                or self.time < 0:
            done = True
            if self.time < 0:
                reward = -100
        else:
            if(self.current_position == self.end_position):
                done = True
                resolved = True
                reward = self.energy * 2
            reward += (self.energy - self.energy_history) / 100
            self.energy_history = self.energy

        return done, resolved, reward

    def get_score(self):
        return self.energy / self.battery_capacity

    def step(self, action):
        graph = self.route_info.graph
        start = self.current_position
        next = action

        if not self.is_action_available(next):
            return self.get_observation(), -0.5, True, False, {}

        #print('start = {}, next = {}, end = {}'.format(start, next, self.end_position))
        consumption_rate = self.consumption_rate / 100  # 耗电量kWh

        route = graph.edges[(start, next)]

        energy_consumed = route['Distance'] * consumption_rate
        travel_time = route['Time']

        if route['IsCharging'] == 1:
            energy_charged = self.charging_power * travel_time  # 充电量kWh
        else:
            energy_charged = 0

        net_energy = energy_consumed - energy_charged  # 净能耗kWh

        self.current_position = next
        self.energy = self.energy - net_energy
        self.time = self.time - travel_time

        self.path.append(self.current_position)
        done, resolved, reward = self.get_reward()
        return self.get_observation(), reward, done, resolved, {}

    def is_action_available(self, action):
        if action < self.action_low or action > self.action_high or self.current_position == action:
            return False
        return True

    def get_observation(self):
        return np.array([self.current_position,
                         self.time,
                         self.energy], dtype=np.float32)
    
    def close(self):
        pass


if __name__ == "__main__":
    pass